import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
from PIL import Image
import numpy as np

class ResNet152FeatureExtractor:
    def __init__(self, device='cuda'):
        # 加载预训练的ResNet152模型
        self.model = models.resnet152(pretrained=True)
        
        # 移除最后的全连接层，只保留特征提取部分
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        
        # 设置为评估模式并移动到指定设备
        self.model.eval()
        self.device = device
        self.model = self.model.to(device)
        
        # 定义图像预处理转换
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

    def extract_features(self, image_path):
        """
        从单张图像中提取特征
        
        Args:
            image_path (str): 图像文件的路径
            
        Returns:
            np.ndarray: 2048维的特征向量
        """
        # 加载并预处理图像
        image = Image.open(image_path).convert('RGB')
        image = self.transform(image).unsqueeze(0)
        image = image.to(self.device)
        
        # 提取特征
        with torch.no_grad():
            features = self.model(image)
            
        # 将特征转换为numpy数组并展平
        features = features.squeeze().cpu().numpy()
        return features

    def extract_features_batch(self, image_paths, batch_size=32):
        """
        批量提取多张图像的特征
        
        Args:
            image_paths (list): 图像文件路径列表
            batch_size (int): 批处理大小
            
        Returns:
            np.ndarray: 特征矩阵，形状为 (n_images, 2048)
        """
        all_features = []
        
        for i in range(0, len(image_paths), batch_size):
            batch_paths = image_paths[i:i + batch_size]
            batch_images = []
            
            # 预处理批次中的所有图像
            for path in batch_paths:
                image = Image.open(path).convert('RGB')
                image = self.transform(image)
                batch_images.append(image)
            
            # 将批次转换为张量
            batch_tensor = torch.stack(batch_images).to(self.device)
            
            # 提取特征
            with torch.no_grad():
                batch_features = self.model(batch_tensor)
                
            # 将特征添加到列表中
            batch_features = batch_features.squeeze().cpu().numpy()
            if len(batch_paths) == 1:
                batch_features = batch_features.reshape(1, -1)
            all_features.append(batch_features)
        
        # 合并所有批次的特征
        return np.vstack(all_features)

def main():
    # 使用示例
    extractor = ResNet152FeatureExtractor()
    
    # 单张图像特征提取示例
    # image_path = "path/to/your/image.jpg"
    # features = extractor.extract_features(image_path)
    # print(f"特征维度: {features.shape}")
    
    # 批量特征提取示例
    # image_paths = ["path1.jpg", "path2.jpg", "path3.jpg"]
    # features = extractor.extract_features_batch(image_paths)
    # print(f"批量特征维度: {features.shape}")

if __name__ == '__main__':
    main() 